import torch
import torch.nn as nn


class RNNLayer(nn.Module):
    def __init__(self, inputs_dim, outputs_dim, recurrent_N, use_orthogonal):
        super(RNNLayer, self).__init__()
        self._recurrent_N = recurrent_N
        self._use_orthogonal = use_orthogonal

        # 使用LSTM代替GRU
        self.rnn = nn.LSTM(inputs_dim, outputs_dim, num_layers=self._recurrent_N)
        for name, param in self.rnn.named_parameters():
            if "bias" in name:
                nn.init.constant_(param, 0)
            elif "weight" in name:
                if self._use_orthogonal:
                    nn.init.orthogonal_(param)
                else:
                    nn.init.xavier_uniform_(param)
        self.norm = nn.LayerNorm(outputs_dim)

    def forward(self, x, hxs, masks):
        # LSTM需要两个隐藏状态：h_0和C_0
        if isinstance(hxs, tuple):
            h_0, c_0 = hxs
        else:
            h_0, c_0 = hxs, None

        if x.size(0) == hxs.size(0):
            # 调整LSTM的输入和隐藏状态
            x, (h_0, c_0) = self.rnn(
                x.unsqueeze(0),
                (h_0 * masks.repeat(1, self._recurrent_N).unsqueeze(-1))
                .transpose(0, 1)
                .contiguous(),
            )
            x = x.squeeze(0)
            h_0 = h_0.transpose(0, 1)
            if c_0 is not None:
                c_0 = c_0.transpose(0, 1)
        else:
            # 处理扁平化的输入
            N = h_0.size(0)
            T = int(x.size(0) / N)

            x = x.view(T, N, x.size(1))
            masks = masks.view(T, N)

            has_zeros = (masks[1:] == 0.0).any(dim=-1).nonzero().squeeze().cpu()
            if has_zeros.dim() == 0:
                has_zeros = [has_zeros.item() + 1]
            else:
                has_zeros = (has_zeros + 1).numpy().tolist()

            has_zeros = [0] + has_zeros + [T]

            h_0 = h_0.transpose(0, 1)

            outputs = []
            for i in range(len(has_zeros) - 1):
                start_idx = has_zeros[i]
                end_idx = has_zeros[i + 1]
                temp = (
                    h_0
                    * masks[start_idx].view(1, -1, 1).repeat(self._recurrent_N, 1, 1)
                ).contiguous()
                rnn_scores, (h_0, c_0) = self.rnn(
                    x[start_idx:end_idx], (temp, c_0[start_idx])
                )
                outputs.append(rnn_scores)

            x = torch.cat(outputs, dim=0)
            x = x.reshape(T * N, -1)
            h_0 = h_0.transpose(0, 1)
            if c_0 is not None:
                c_0 = c_0.transpose(0, 1)

        x = self.norm(x)
        return x, (h_0, c_0)
